package com.zhiyou.app.common.util;

import cn.hutool.core.date.DateUtil;
import lombok.Setter;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.BiConsumer;

/**
 * 增强版定时过期Map，同时支持：
 * 1. 为每个键值对单独设置过期时间
 * 2. 设置统一的默认过期时间
 */
public class EnhancedExpiringMap<K, V> {
    // 存储键值对
    private final Map<K, V> valueMap = new HashMap<>();
    // 存储每个键的过期时间(毫秒时间戳)，null表示使用默认过期时间
    private final Map<K, Long> expiryMap = new HashMap<>();
    // 存储每个键对应的定时任务
    private final Map<K, ScheduledFuture<?>> taskMap = new HashMap<>();

    // 线程池用于执行过期任务
    private final ScheduledExecutorService scheduler;

    // 读写锁保证线程安全
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock();
    private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock();

    /**
     * -- SETTER --
     * 设置过期事件监听器
     */
    // 过期事件监听器
    @Setter
    private BiConsumer<K, V> expiryListener;

    // 默认过期时间和单位
    private long defaultExpiryDuration;
    private TimeUnit defaultExpiryUnit;

    /**
     * 构造方法，使用默认线程池和默认过期时间
     *
     * @param defaultExpiryDuration 默认过期时间长度
     * @param defaultExpiryUnit     默认时间单位
     */
    public EnhancedExpiringMap(long defaultExpiryDuration, TimeUnit defaultExpiryUnit) {
        this(
                defaultExpiryDuration,
                defaultExpiryUnit,
                Executors.newScheduledThreadPool(1, r -> {
                    Thread thread = new Thread(r, "enhanced-expiring-map-scheduler");
                    thread.setDaemon(true);
                    return thread;
                })
        );
    }

    /**
     * 构造方法，允许传入自定义线程池
     */
    public EnhancedExpiringMap(
            long defaultExpiryDuration,
            TimeUnit defaultExpiryUnit,
            ScheduledExecutorService scheduler) {
        if (defaultExpiryDuration <= 0) {
            throw new IllegalArgumentException("默认过期时间必须为正数");
        }
        this.defaultExpiryDuration = defaultExpiryDuration;
        this.defaultExpiryUnit = defaultExpiryUnit;
        this.scheduler = scheduler;
    }

    /**
     * 使用默认过期时间存入键值对
     */
    public void put(K key, V value) {
        put(key, value, defaultExpiryDuration, defaultExpiryUnit, true);
    }

    /**
     * 为特定键值对单独设置过期时间
     */
    public void put(K key, V value, long duration, TimeUnit unit) {
        if (duration <= 0) {
            throw new IllegalArgumentException("过期时间必须为正数");
        }
        put(key, value, duration, unit, false);
    }

    /**
     * 内部实现的put方法
     *
     * @param useDefault 是否使用默认过期时间
     */
    private void put(K key, V value, long duration, TimeUnit unit, boolean useDefault) {
        if (key == null) {
            throw new NullPointerException("key cannot be null");
        }

        long delay = unit.toMillis(duration);
        writeLock.lock();
        try {
            // 取消已有的定时任务
            ScheduledFuture<?> existingTask = taskMap.get(key);
            if (existingTask != null) {
                existingTask.cancel(false);
            }

            // 存储键值对
            valueMap.put(key, value);

            // 存储过期时间标记（null表示使用默认）
            if (useDefault) {
                expiryMap.put(key, null);
            } else {
                expiryMap.put(key, System.currentTimeMillis() + delay);
            }

            // 创建新的定时任务
            ScheduledFuture<?> task = scheduler.schedule(() -> expireKey(key), delay, TimeUnit.MILLISECONDS);
            taskMap.put(key, task);
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * 获取键对应的值，如果已过期则返回null
     */
    public V get(K key) {
        readLock.lock();
        try {
            // 检查是否存在
            if (!valueMap.containsKey(key)) {
                return null;
            }

            // 检查是否已过期
            Long expiryTime = expiryMap.get(key);
            long now = System.currentTimeMillis();

            // 如果是默认过期时间，计算当前应该过期的时间
            if (expiryTime == null) {
                // 获取任务创建时间（通过计算反推）
                ScheduledFuture<?> task = taskMap.get(key);
                if (task == null) {
                    return null;
                }

                // 对于默认过期时间，判断任务是否已完成
                if (task.isDone()) {
                    return null;
                }
            }
            // 对于自定义过期时间，直接比较时间戳
            else if (now > expiryTime) {
                return null;
            }

            return valueMap.get(key);
        } finally {
            readLock.unlock();
        }
    }

    /**
     * 更新默认过期时间（只影响新添加的键值对）
     */
    public void updateDefaultExpiry(long duration, TimeUnit unit) {
        if (duration <= 0) {
            throw new IllegalArgumentException("过期时间必须为正数");
        }
        writeLock.lock();
        try {
            this.defaultExpiryDuration = duration;
            this.defaultExpiryUnit = unit;
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * 移除键并取消对应的过期任务
     */
    public V remove(K key) {
        writeLock.lock();
        try {
            // 取消定时任务
            ScheduledFuture<?> task = taskMap.remove(key);
            if (task != null) {
                task.cancel(false);
            }
            // 移除过期时间记录
            expiryMap.remove(key);
            // 返回并移除值
            return valueMap.remove(key);
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * 检查键是否存在且未过期
     */
    public boolean containsKey(K key) {
        return get(key) != null;
    }

    /**
     * 获取当前未过期的键集合
     */
    public Set<K> keySet() {
        readLock.lock();
        try {
            Set<K> activeKeys = new HashSet<>();
            long now = System.currentTimeMillis();

            for (K key : valueMap.keySet()) {
                Long expiryTime = expiryMap.get(key);

                if (expiryTime == null) {
                    // 默认过期时间，检查任务是否已完成
                    ScheduledFuture<?> task = taskMap.get(key);
                    if (task != null && !task.isDone()) {
                        activeKeys.add(key);
                    }
                } else if (now <= expiryTime) {
                    // 自定义过期时间且未过期
                    activeKeys.add(key);
                }
            }

            return activeKeys;
        } finally {
            readLock.unlock();
        }
    }

    /**
     * 获取当前未过期的键值对数量
     */
    public int size() {
        return keySet().size();
    }

    /**
     * 清空Map并取消所有定时任务
     */
    public void clear() {
        writeLock.lock();
        try {
            // 取消所有定时任务
            for (ScheduledFuture<?> task : taskMap.values()) {
                task.cancel(false);
            }
            taskMap.clear();
            expiryMap.clear();
            valueMap.clear();
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * 处理键过期逻辑
     */
    private void expireKey(K key) {
        writeLock.lock();
        try {
            // 移除键值对
            V value = valueMap.remove(key);
            expiryMap.remove(key);
            taskMap.remove(key);

            // 触发过期事件
            if (expiryListener != null && value != null) {
                expiryListener.accept(key, value);
            }
        } finally {
            writeLock.unlock();
        }
    }

    /**
     * 关闭线程池，释放资源
     */
    public void shutdown() {
        scheduler.shutdown();
    }

    public static void main(String[] args) {
        System.out.println("start：" + DateUtil.format(new Date(), "HH:mm:ss"));
        EnhancedExpiringMap<String, String> map = new EnhancedExpiringMap<>(30, TimeUnit.SECONDS);
        map.setExpiryListener((k, v) -> {
            System.out.println(DateUtil.format(new Date(), "HH:mm:ss") + "过期：" + "key:" + k + " value:" + v);
        });
        map.put("key1", "value1", 5, TimeUnit.SECONDS);
        map.put("key2", "value2", 8, TimeUnit.SECONDS);
        map.put("key3", "value3");
        new Thread(() -> {
            try {
                TimeUnit.SECONDS.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }).start();
    }
}

